"""
Pro-DUCA (Proximal Dual Consensus Algorithm) 算法实现

DUCA的近端变体，适用于无界可行域
"""

import numpy as np
from typing import Dict, Optional
from .base_algorithm import BaseDistributedAlgorithm
import cvxpy as cp


class ProDUCA(BaseDistributedAlgorithm):
    """Pro-DUCA算法实现"""
    
    def __init__(self, problem, network, **kwargs):
        """
        初始化Pro-DUCA算法
        
        参数:
            problem: 分布式优化问题实例
            network: 网络拓扑实例
            rho: 近端参数（默认1.0）
            tau: 对偶近端参数（默认1.0）
            mode: 通信模式 ('single-exchange' 或 'double-exchange')
            step_size: 对偶变量步长
        """
        super().__init__(problem, network, **kwargs)
        
        # Pro-DUCA特定参数
        self.rho = kwargs.get('rho', 1.0)
        self.tau = kwargs.get('tau', 1.0)
        self.mode = kwargs.get('mode', 'single-exchange')
        self.step_size = kwargs.get('step_size', None)
        
        # 权重矩阵
        self.W = network.weight_matrix
        
        # 自适应步长
        if self.step_size is None:
            eigenvalues = network.get_weight_eigenvalues()
            self.step_size = 0.5 / (1 + abs(eigenvalues[1]))
            
        # 近端中心点
        self.x_center = {}
        self.y_center = {}
        
    def _initialize_auxiliary_variables(self):
        """初始化Pro-DUCA的辅助变量"""
        # 初始化近端中心点
        for i in range(self.num_nodes):
            self.x_center[i] = np.zeros(self.problem.dimensions[i])
            self.y_center[i] = np.zeros(self.m + self.p)
            self.z[i] = np.zeros(self.m + self.p)
            
        # 双交换模式的额外变量
        if self.mode == 'double-exchange':
            self.v = {}
            for i in range(self.num_nodes):
                self.v[i] = np.zeros(self.m + self.p)
                
    def _update_primal_variables(self):
        """更新原始变量（带近端项）"""
        for i in range(self.num_nodes):
            # 解决带近端项的局部子问题
            # minimize f_i(x_i) + <y_i, G_i(x_i)> + (rho/2)||x_i - x_center_i||^2
            # subject to x_i ∈ X_i
            
            x_i = cp.Variable(self.problem.dimensions[i])
            
            # 目标函数
            obj = self.problem.objectives[i](x_i)
            
            # 对偶项
            lambda_i = self.y[i][:self.m]
            nu_i = self.y[i][self.m:]
            
            dual_term = 0
            
            # 不等式约束贡献
            for j, g_ij in enumerate(self.problem.inequalities[i]):
                if j < self.m:
                    dual_term += lambda_i[j] * g_ij(x_i)
                    
            # 等式约束贡献
            for j, h_ij in enumerate(self.problem.equalities[i]):
                if j < self.p:
                    dual_term += nu_i[j] * h_ij(x_i)
                    
            # 近端项
            proximal_term = (self.rho / 2) * cp.sum_squares(x_i - self.x_center[i])
            
            # 总目标
            objective = obj + dual_term + proximal_term
            
            # 局部约束
            constraints = []
            if self.problem.local_sets[i] is not None:
                set_type = self.problem.local_sets[i]['type']
                if set_type == 'box':
                    lower = self.problem.local_sets[i]['lower']
                    upper = self.problem.local_sets[i]['upper']
                    constraints.append(x_i >= lower)
                    constraints.append(x_i <= upper)
                elif set_type == 'ball':
                    center = self.problem.local_sets[i]['center']
                    radius = self.problem.local_sets[i]['radius']
                    constraints.append(cp.norm(x_i - center) <= radius)
                elif set_type == 'simplex':
                    constraints.append(x_i >= 0)
                    constraints.append(cp.sum(x_i) <= 1)
                    
            # 求解子问题
            try:
                prob = cp.Problem(cp.Minimize(objective), constraints)
                prob.solve(solver=cp.OSQP, verbose=False)
                
                if prob.status == cp.OPTIMAL:
                    self.x[i] = x_i.value
                else:
                    # 使用近端梯度法作为备用
                    self._proximal_gradient_step(i)
            except:
                self._proximal_gradient_step(i)
                
    def _proximal_gradient_step(self, i: int):
        """近端梯度步（备用求解器）"""
        # 计算梯度
        grad = self.problem.objectives[i].gradient(self.x[i])
        
        # 添加对偶项梯度
        lambda_i = self.y[i][:self.m]
        nu_i = self.y[i][self.m:]
        
        for j, g_ij in enumerate(self.problem.inequalities[i]):
            if j < self.m and hasattr(g_ij, 'gradient'):
                grad += lambda_i[j] * g_ij.gradient(self.x[i])
                
        for j, h_ij in enumerate(self.problem.equalities[i]):
            if j < self.p and hasattr(h_ij, 'gradient'):
                grad += nu_i[j] * h_ij.gradient(self.x[i])
                
        # 添加近端项梯度
        grad += self.rho * (self.x[i] - self.x_center[i])
        
        # 梯度步
        step_size = 0.1 / (self.rho + np.linalg.norm(grad))
        x_new = self.x[i] - step_size * grad
        
        # 投影
        if self.problem.local_sets[i] is not None:
            self.x[i] = self.problem.project_to_local_set(x_new, i)
        else:
            self.x[i] = x_new
            
    def _update_dual_variables(self):
        """更新对偶变量（带近端项）"""
        if self.mode == 'single-exchange':
            self._single_exchange_dual_update_proximal()
        else:
            self._double_exchange_dual_update_proximal()
            
    def _single_exchange_dual_update_proximal(self):
        """单交换模式的近端对偶更新"""
        # 计算约束值
        g, h = self.problem.evaluate_constraints(self.x)
        c = np.concatenate([g, h])
        
        y_new = {}
        for i in range(self.num_nodes):
            # 共识步骤
            y_consensus = np.zeros(self.m + self.p)
            for j in range(self.num_nodes):
                y_consensus += self.W[i, j] * self.y[j]
                
            # 近端梯度上升
            y_tilde = y_consensus + self.step_size * c / self.num_nodes
            
            # 近端算子
            y_new[i] = (self.tau * self.y_center[i] + y_tilde) / (1 + self.tau)
            
            # 投影
            y_new[i][:self.m] = np.maximum(y_new[i][:self.m], 0)
            
        self.y = y_new
        
    def _double_exchange_dual_update_proximal(self):
        """双交换模式的近端对偶更新"""
        # 第一次交换
        v_new = {}
        for i in range(self.num_nodes):
            v_new[i] = np.zeros(self.m + self.p)
            for j in range(self.num_nodes):
                v_new[i] += self.W[i, j] * self.v[j]
        self.v = v_new
        
        # 计算约束
        g, h = self.problem.evaluate_constraints(self.x)
        c = np.concatenate([g, h])
        
        # 更新对偶变量
        y_new = {}
        for i in range(self.num_nodes):
            # 梯度上升
            y_tilde = self.v[i] + self.step_size * c / self.num_nodes
            
            # 近端算子
            y_new[i] = (self.tau * self.y_center[i] + y_tilde) / (1 + self.tau)
            
            # 投影
            y_new[i][:self.m] = np.maximum(y_new[i][:self.m], 0)
            
        self.y = y_new
        
        # 第二次交换
        v_new = {}
        for i in range(self.num_nodes):
            v_new[i] = np.zeros(self.m + self.p)
            for j in range(self.num_nodes):
                v_new[i] += self.W[i, j] * self.y[j]
        self.v = v_new
        
    def _update_auxiliary_variables(self):
        """更新辅助变量（包括近端中心点）"""
        # 更新原始变量的近端中心
        for i in range(self.num_nodes):
            self.x_center[i] = self.x[i].copy()
            
        # 更新对偶变量的近端中心
        for i in range(self.num_nodes):
            self.y_center[i] = self.y[i].copy()
            
        # 更新全局对偶估计
        for i in range(self.num_nodes):
            self.z[i] = self.y[i].copy()
            
    def update_proximal_parameters(self, rho: Optional[float] = None, 
                                  tau: Optional[float] = None):
        """动态更新近端参数"""
        if rho is not None:
            self.rho = rho
        if tau is not None:
            self.tau = tau
            
    def get_proximal_residual(self) -> float:
        """计算近端残差"""
        prox_res = 0.0
        
        # 原始变量近端残差
        for i in range(self.num_nodes):
            prox_res += np.linalg.norm(self.x[i] - self.x_center[i])**2
            
        # 对偶变量近端残差
        for i in range(self.num_nodes):
            prox_res += np.linalg.norm(self.y[i] - self.y_center[i])**2
            
        return np.sqrt(prox_res) 